import os
import re
import sys
from PIL import Image
from ebooklib import epub
from datetime import datetime

def get_images(subfolder, extensions=None):
    """获取指定文件夹中的所有图片文件，并按文件名排序"""
    if extensions is None:
        extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff'}

    images = []
    for root, _, files in os.walk(subfolder):
        for file in files:
            if any(file.lower().endswith(ext) for ext in extensions):
                images.append(os.path.join(root, file))

    return sorted(images, key=custom_sort)

def custom_sort(s):
    """数字感知的排序函数"""
    filename = os.path.basename(s)
    if "-" in filename:
        filename = filename.split("-")[-1]
    numbers = re.findall(r'\d+', filename)
    if numbers:
        return int(''.join(numbers))
    return 0

def create_epub(folder_path):
    """将图片文件夹转换为EPUB格式"""
    try:
        # 创建EPUB书籍
        book = epub.EpubBook()
        
        # 设置元数据
        folder_name = os.path.basename(folder_path)
        book.set_identifier(f'id_{folder_name}_{datetime.now().strftime("%Y%m%d%H%M%S")}')
        book.set_title(folder_name)
        book.set_language('zh-CN')
        book.add_author('IMG2EPUB Generator')

        # 获取所有图片
        images = get_images(folder_path)
        if not images:
            print("错误：未找到任何图片文件")
            return False

        # 创建章节和目录
        chapters = []
        toc = []
        spine = ['nav']

        print(f"\n开始处理 {len(images)} 张图片...")

        # 处理每张图片
        for index, image_path in enumerate(images, 1):
            try:
                # 创建图片章节
                chapter_title = f'Page {index}'
                chapter = epub.EpubHtml(
                    title=chapter_title,
                    file_name=f'page_{index}.xhtml'
                )

                # 添加图片到EPUB
                img_filename = os.path.basename(image_path)
                epub_img_name = f'images/{index}_{img_filename}'
                
                # 读取并添加图片
                with open(image_path, 'rb') as img_file:
                    epub_image = epub.EpubItem(
                        uid=f'image_{index}',
                        file_name=epub_img_name,
                        media_type='image/jpeg',
                        content=img_file.read()
                    )
                book.add_item(epub_image)

                # 创建包含图片的HTML内容
                chapter.content = f'''<?xml version="1.0" encoding="utf-8"?>
                <!DOCTYPE html>
                <html xmlns="http://www.w3.org/1999/xhtml" xmlns:epub="http://www.idpf.org/2007/ops">
                <head>
                    <title>{chapter_title}</title>
                    <style>
                        img {{ max-width: 100%; height: auto; display: block; margin: 0 auto; }}
                        body {{ margin: 0; padding: 0; }}
                    </style>
                </head>
                <body>
                    <img src="{epub_img_name}" alt="{chapter_title}" />
                </body>
                </html>'''

                # 添加章节
                book.add_item(chapter)
                chapters.append(chapter)
                toc.append(epub.Link(f'page_{index}.xhtml', chapter_title, f'page_{index}'))
                spine.append(chapter)

                print(f"处理进度: [{index}/{len(images)}] {img_filename}")

            except Exception as e:
                print(f"警告：处理图片 {img_filename} 时出错: {str(e)}")
                continue

        # 添加导航
        book.toc = toc
        book.spine = spine
        book.add_item(epub.EpubNcx())
        book.add_item(epub.EpubNav())

        # 生成EPUB文件
        epub_path = os.path.join(os.path.dirname(folder_path), f"{folder_name}.epub")
        epub.write_epub(epub_path, book)

        print(f"\n转换完成！")
        print(f"EPUB文件已保存至：{epub_path}")
        return True

    except Exception as e:
        print(f"错误：转换过程中出现异常: {str(e)}")
        return False

def main():
    """主函数"""
    print("\n=== 图片文件夹转EPUB工具 ===")
    
    if len(sys.argv) > 1:
        folder_path = sys.argv[1]
    else:
        folder_path = input("请输入图片文件夹路径: ").strip()
    
    # 处理路径中的引号
    folder_path = folder_path.strip('"\'')
    
    # 验证路径
    if not os.path.exists(folder_path):
        print(f"错误: 文件夹 '{folder_path}' 不存在")
        return
    
    if not os.path.isdir(folder_path):
        print(f"错误: '{folder_path}' 不是一个文件夹")
        return
        
    # 开始转换
    create_epub(folder_path)

if __name__ == '__main__':
    main()
